import os
import sys
import argparse
from pathlib import Path
import traceback
from pdf2image import convert_from_path
import pytesseract
from tqdm import tqdm

def ocr_pdf_to_txt(pdf_path):
    """
    使用OCR技术将PDF转换为TXT文件
    适合处理扫描件PDF
    
    参数:
    pdf_path (str): PDF文件的路径
    
    返回:
    str: 创建的TXT文件的路径
    """
    # 创建与原文件同名但扩展名为.txt的输出文件路径
    output_path = os.path.splitext(pdf_path)[0] + '.txt'
    
    try:
        print(f"  开始OCR处理文件: {pdf_path}")
        
        # 将PDF转换为图像
        print(f"  将PDF转换为图像...")
        images = convert_from_path(pdf_path, dpi=300)
        
        if not images:
            print(f"  错误: 无法将PDF转换为图像 - {pdf_path}")
            return None
        
        print(f"  转换成功，共 {len(images)} 页")
        
        # 对每个图像进行OCR处理
        text = ""
        for i, image in enumerate(tqdm(images, desc="OCR处理页面", unit="页")):
            print(f"  处理第 {i+1}/{len(images)} 页...")
            # 使用pytesseract进行OCR
            page_text = pytesseract.image_to_string(image, lang='chi_sim+eng')
            if page_text:
                text += page_text + "\n\n"
            else:
                print(f"  警告: 页面 {i+1} 未识别到文本")
        
        if not text.strip():
            print(f"  警告: 未能从PDF中提取到任何文本 - {pdf_path}")
            return None
        
        # 写入TXT文件
        with open(output_path, 'w', encoding='utf-8') as txt_file:
            txt_file.write(text)
        
        print(f"  成功: {pdf_path} -> {output_path}")
        return output_path
        
    except Exception as e:
        print(f"  错误转换 {pdf_path}: {str(e)}")
        print(f"  详细错误: {traceback.format_exc()}")
        return None

def process_specific_pdfs(pdf_files):
    """
    处理指定的PDF文件列表
    
    参数:
    pdf_files (list): PDF文件路径列表
    
    返回:
    tuple: (成功转换的文件数, 失败转换的文件数)
    """
    success_count = 0
    failure_count = 0
    
    total_files = len(pdf_files)
    print(f"准备处理 {total_files} 个PDF文件:")
    
    for i, pdf_path in enumerate(pdf_files):
        pdf_path = Path(pdf_path)
        if not pdf_path.exists():
            print(f"错误: 文件不存在 - {pdf_path}")
            failure_count += 1
            continue
            
        if not pdf_path.is_file() or pdf_path.suffix.lower() != '.pdf':
            print(f"错误: 不是有效的PDF文件 - {pdf_path}")
            failure_count += 1
            continue
            
        print(f"\n处理文件 {i+1}/{total_files}: {pdf_path.name}")
        
        # 使用OCR处理PDF
        result = ocr_pdf_to_txt(str(pdf_path))
        
        if result:
            success_count += 1
        else:
            failure_count += 1
    
    return success_count, failure_count

def main():
    # 创建命令行参数解析器
    parser = argparse.ArgumentParser(description='使用OCR技术将扫描版PDF转换为TXT文件')
    parser.add_argument('pdf_files', nargs='+', help='要处理的PDF文件路径')
    parser.add_argument('--lang', default='chi_sim+eng', help='OCR语言，默认为中文简体+英文')
    
    # 解析命令行参数
    args = parser.parse_args()
    
    # 检查tesseract是否已安装
    try:
        pytesseract.get_tesseract_version()
    except pytesseract.TesseractNotFoundError:
        print("错误: 未找到Tesseract OCR。请先安装Tesseract OCR。")
        print("Ubuntu/Debian: sudo apt-get install tesseract-ocr")
        print("MacOS: brew install tesseract")
        print("Windows: 下载并安装 https://github.com/UB-Mannheim/tesseract/wiki")
        print("同时也需要安装相应的语言包，如中文: sudo apt-get install tesseract-ocr-chi-sim")
        sys.exit(1)
    
    # 设置tesseract语言
    pytesseract.pytesseract.tesseract_cmd = r'tesseract'
    
    # 打印系统信息
    print(f"当前工作目录: {os.getcwd()}")
    print(f"Python 版本: {sys.version}")
    print(f"Tesseract 版本: {pytesseract.get_tesseract_version()}")
    print(f"OCR语言: {args.lang}")
    print(f"处理文件数: {len(args.pdf_files)}\n")
    
    # 执行OCR处理
    success, failure = process_specific_pdfs(args.pdf_files)
    
    # 打印转换结果摘要
    print(f"\n转换完成：")
    print(f"成功转换: {success} 文件")
    print(f"失败转换: {failure} 文件")

if __name__ == "__main__":
    main()